#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2021/9/14 12:10 上午
# @Author  : khan_long
# @Email   : longkehan15@qq.com
# @File    : dataset_reader.py
# @Software: PyCharm

import jsonlines
import re
import numpy as np
import spacy

regex_find_citation = re.compile(r"\(\s?(([A-Za-z\-]+\s)+([A-Za-z\-]+\.?)?,?\s\d{2,4}[a-c]?(;\s)?)+\s?\)|"
                                 r"\[(\d{1,3},\s?)+\d{1,3}\]|"
                                 r"\[[\d,-]+\]|(\([A-Z][a-z]+, \d+[a-c]?\))|"
                                 r"([A-Z][a-z]+ (et al\.)? \(\d+[a-c]?\))|"
                                 r"[A-Z][a-z]+ and [A-Z][a-z]+ \(\d+[a-c]?\)]")

spacy_tokenizer = spacy.load('en_core_web_sm')


def my_tokenizer(_sent):
    # TODO remove punctuation
    return [token.text for token in spacy_tokenizer(_sent)]


def section_title_dataset_reader(_which='acl', _clean_citation=False, _tokenizer=my_tokenizer):
    """
    :param _tokenizer:
    :param _which: aclarc or scicite, 'acl' means the former, otherwise the later
    :param _clean_citation:
    :return: List[(List(str), str)]
    """
    examples = []
    if _which == 'acl':
        _path = './datasets/acl-arc/scaffolds/sections-scaffold-train.jsonl'
    else:
        _path = './datasets/scicite/scaffolds/sections-scaffold-train.jsonl'
    for obj in jsonlines.open(_path):
        citation_text = obj['text']
        if _clean_citation:
            citation_text = regex_find_citation.sub("", citation_text)
        section_name = obj['section_name']
        examples.append((_tokenizer(citation_text), section_name))
    return examples


def citation_worthiness_reader(_which='acl', _clean_citation=True, _tokenizer=my_tokenizer):
    """
    :param _tokenizer:
    :param _which: aclarc or scicite, 'acl' means the former, otherwise the later
    :param _clean_citation:
    :return: List[(List(str),bool)]
    """
    examples = []
    if _which == 'acl':
        _path = './datasets/acl-arc/scaffolds/cite-worthiness-scaffold-train.jsonl'
    else:
        _path = './datasets/scicite/scaffolds/cite-worthiness-scaffold-train.jsonl'
    for obj in jsonlines.open(_path):
        if _clean_citation:
            citation_text = obj['cleaned_cite_text']
        else:
            citation_text = obj['text']
        is_citation = obj['is_citation']
        examples.append((_tokenizer(citation_text), is_citation))
    return examples


def aan_cc_dataset_reader(_version='aan_2013_50_2'):
    train_examples = []
    for obj in jsonlines.open('./datasets/aan/' + _version + '/train.jsonl'):
        pass


if __name__ == '__main__':
    s_examples_acl = section_title_dataset_reader(_which='acl', _tokenizer=lambda x: x.split())
    s_examples_sci = section_title_dataset_reader(_which='sci', _tokenizer=lambda x: x.split())
    print(set([one[1] for one in s_examples_acl]))
    print(set([one[1] for one in s_examples_sci]))
    w_examples_acl = citation_worthiness_reader(_which='acl')
    w_examples_sci = citation_worthiness_reader(_which='sci')
    print(1)
